#pragma once

#include <iostream>
#include <vector>

#include "activation.h"
#include "linearlayer.h"

template <typename T>
class Conv2d
{
private:
    Tensor<T> kernel;    // 卷积核：[out_channels, in_channels, k_h, k_w]
    Tensor<T> bias;      // 偏置：[out_channels]
    int in_channels;     // 输入通道数
    int out_channels;    // 输出通道数
    int kernel_h;        // 卷积核高度
    int kernel_w;        // 卷积核宽度
    int stride;          // 步长
    int padding;         // 填充

    int compute_out_size(int in_size) const {
        return (in_size + 2 * padding - kernel_h) / stride + 1;
    }

    // 对输入进行填充（在边缘补0）
    Tensor<T> pad_input(const Tensor<T>& input) const {
        if (padding == 0) return input;  // 无需填充直接返回

        // 输入形状：[batch, in_channels, h, w]
        const auto& in_shape = input.get_shape();
        int batch = in_shape[0];
        int c = in_shape[1];
        int h = in_shape[2];
        int w = in_shape[3];

        // 填充后形状：[batch, c, h+2*padding, w+2*padding]
        Tensor<T> padded({ batch, c, h + 2 * padding, w + 2 * padding });
        padded.fill(T{ 0 });  // 初始化为0

        // 复制原始数据到填充后的中心区域
        auto& input_data = input.get_data();
        auto& padded_data = padded.get_data();
        int in_stride_h = w;                  // 输入的高度方向步长（每行元素数）
        int in_stride_c = h * w;              // 输入的通道方向步长（每个通道元素数）
        int in_stride_b = c * h * w;          // 输入的batch方向步长（每个样本元素数）
        int pad_stride_h = w + 2 * padding;     // 填充后的高度方向步长
        int pad_stride_c = (h + 2 * padding) * (w + 2 * padding);  // 填充后的通道步长
        int pad_stride_b = c * pad_stride_c;  // 填充后的batch步长

        for (int b = 0; b < batch; ++b) {
            for (int ch = 0; ch < c; ++ch) {
                for (int i = 0; i < h; ++i) {
                    for (int j = 0; j < w; ++j) {
                        // 输入的线性索引
                        int in_idx = b * in_stride_b + ch * in_stride_c + i * in_stride_h + j;
                        // 填充后张量的对应索引（偏移padding）
                        int pad_idx = b * pad_stride_b + ch * pad_stride_c + (i + padding) * pad_stride_h + (j + padding);
                        padded_data[pad_idx] = input_data[in_idx];
                    }
                }
            }
        }
        return padded;
    }

public:
    Conv2d(int in_channels, int out_channels, int kernel_size, int stride = 1, int padding = 0)
        : in_channels(in_channels), out_channels(out_channels),
        kernel_h(kernel_size), kernel_w(kernel_size),  // 简化为正方形kernel
        stride(stride), padding(padding) {
        // 初始化卷积核：[out_channels, in_channels, kernel_h, kernel_w]
        kernel = Tensor<T>({ out_channels, in_channels, kernel_h, kernel_w });
        kernel.random(-0.1, 0.1);  // 简单随机初始化

        // 初始化偏置：[out_channels]
        bias = Tensor<T>::zeros({ out_channels });
    }

    // 前向传播：输入形状[batch, in_channels, h, w] → 输出[batch, out_channels, out_h, out_w]
    Tensor<T> forward(const Tensor<T>& input) {
        // 检查输入形状是否符合要求（必须是4D：[batch, c, h, w]）
        if (input.ndim() != 4) {
            throw std::invalid_argument("Conv2d input must be 4D (batch, in_channels, h, w)");
        }
        const auto& in_shape = input.get_shape();
        int batch = in_shape[0];
        int in_c = in_shape[1];
        int in_h = in_shape[2];
        int in_w = in_shape[3];
        if (in_c != in_channels) {
            throw std::invalid_argument("Input channels do not match Conv2d in_channels");
        }

        // 1. 对输入进行填充
        Tensor<T> padded = pad_input(input);
        const auto& pad_shape = padded.get_shape();
        int pad_h = pad_shape[2];
        int pad_w = pad_shape[3];

        // 2. 计算输出特征图尺寸
        int out_h = compute_out_size(in_h);
        int out_w = compute_out_size(in_w);

        // 3. 初始化输出张量：[batch, out_channels, out_h, out_w]
        Tensor<T> output({ batch, out_channels, out_h, out_w });
        output.fill(T{ 0 });
        auto& output_data = output.get_data();
        auto& kernel_data = kernel.get_data();
        auto& bias_data = bias.get_data();

        // 4. 计算卷积（核心逻辑）
        // 输出张量的步长（用于计算线性索引）
        int out_stride_oc = out_h * out_w;                // 输出通道方向步长
        int out_stride_b = out_channels * out_stride_oc;  // batch方向步长

        // 卷积核的步长
        int kernel_stride_ic = kernel_h * kernel_w;              // 输入通道方向步长
        int kernel_stride_oc = in_channels * kernel_stride_ic;   // 输出通道方向步长

        // 填充后输入的步长
        int pad_stride_c = pad_h * pad_w;              // 通道方向步长
        int pad_stride_b = in_channels * pad_stride_c; // batch方向步长

        // 遍历每个batch
        for (int b = 0; b < batch; ++b) {
            // 遍历每个输出通道
            for (int oc = 0; oc < out_channels; ++oc) {
                // 遍历输出特征图的每个像素
                for (int i = 0; i < out_h; ++i) {
                    for (int j = 0; j < out_w; ++j) {
                        // 计算当前输出像素的线性索引
                        int out_idx = b * out_stride_b + oc * out_stride_oc + i * out_w + j;

                        // 累加所有输入通道的卷积结果
                        T sum = T{ 0 };
                        for (int ic = 0; ic < in_channels; ++ic) {
                            // 遍历卷积核的每个元素
                            for (int ki = 0; ki < kernel_h; ++ki) {
                                for (int kj = 0; kj < kernel_w; ++kj) {
                                    // 卷积核当前元素的索引
                                    int kernel_idx = oc * kernel_stride_oc + ic * kernel_stride_ic + ki * kernel_w + kj;
                                    T k_val = kernel_data[kernel_idx];

                                    // 输入特征图上对应位置的索引（滑动窗口）
                                    int in_i = i * stride + ki;  // 步长*当前输出行 + kernel行偏移
                                    int in_j = j * stride + kj;  // 步长*当前输出列 + kernel列偏移
                                    int pad_idx = b * pad_stride_b + ic * pad_stride_c + in_i * pad_w + in_j;
                                    T x_val = padded.get_data()[pad_idx];

                                    sum += k_val * x_val;  // 内积累加
                                }
                            }
                        }

                        // 加上偏置，存入输出
                        output_data[out_idx] = sum + bias_data[oc];
                    }
                }
            }
        }

        return output;
    }

    const Tensor<T>& get_kernel() const { return kernel; }
    const Tensor<T>& get_bias() const { return bias; }
};

